#!/usr/bin/env python
import numpy as np
from MapUtility import loadMap

def LogPoisson(model, data, mask0=None, alldata=True):
    mask = model > 0.
    if not mask0 is None:
        mask = mask * mask0>0.

    if alldata:
        model2, data2 = model[mask], data[mask]
        return 2.*(model2 - data2*np.log(model2)).sum()
    else:
        outdata = []
        for ilayer in xrange(data.shape[0]):
            mask2 = mask[ilayer, ...]
            model2, data2 = model[ilayer, ...][mask2], data[ilayer, ...][mask2]
            outdata.append(2.*(model2 - data2*np.log(model2)).sum())
        return np.atleast_1d(outdata)

def loglike(ccubefile, mcubefile, maskfile=None, ccubeExtIdx=None, mcubeExtIdx=None, maskExtIdx=None, alldata=True):
    ccube = loadMap(ccubefile, extIdx=ccubeExtIdx)
    mcube = loadMap(mcubefile, extIdx=mcubeExtIdx)
    if not maskfile is None:
        mask = loadMap(maskfile, extIdx=maskExtIdx)
    else:
        mask = None
    return LogPoisson(mcube, ccube, mask, alldata)

def myNone(inputs):
    if isinstance(inputs, str) and inputs.lower() == 'none':
        return None
    else:
        return inputs

def cli():
    import argparse

    parser = argparse.ArgumentParser(description='Calculate loglike given the mcube and ccube.')
    parser.add_argument("ccubefile", type=str, help='ccube')
    parser.add_argument("mcubefile", type=str, help='mcube')
    parser.add_argument("-maskfile", type=myNone, default=None, help='mask')
    parser.add_argument("-ccubeExtIdx", type=myNone, default=None, help='the extension index of the ccube')
    parser.add_argument("-mcubeExtIdx", type=myNone, default=None, help='the extension index of the mcube')
    parser.add_argument("-maskExtIdx", type=myNone, default=None, help='the extension index of the mask')
    parser.add_argument("-e", action='store_true', default=False, help='show the value in each energy bin')
    args = parser.parse_args()

    outdata = loglike(ccubefile=args.ccubefile, mcubefile=args.mcubefile, maskfile=args.maskfile,
                      ccubeExtIdx=args.ccubeExtIdx, mcubeExtIdx=args.mcubeExtIdx, maskExtIdx=args.maskExtIdx,
                      alldata=(not args.e))
    if outdata.ndim == 0:
        print 'logLike: %s' % outdata
    else:
        for ival, val in enumerate(outdata):
            print 'logLike[%s]: %s' % (ival, val)

if __name__ == '__main__': cli()
